/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.solr.cloud;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import org.apache.solr.client.solrj.cloud.DistribStateManager;
import org.apache.solr.client.solrj.cloud.SolrCloudManager;
import org.apache.solr.client.solrj.impl.CloudSolrClient;
import org.apache.solr.client.solrj.request.CollectionAdminRequest;
import org.apache.solr.client.solrj.response.CollectionAdminResponse;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.client.solrj.response.RequestStatusState;
import org.apache.solr.cloud.api.collections.ReindexCollectionCmd;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.common.cloud.ClusterState;
import org.apache.solr.common.cloud.DocCollection;
import org.apache.solr.common.cloud.ImplicitDocRouter;
import org.apache.solr.common.cloud.ZkStateReader;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.util.TimeSource;
import org.apache.solr.util.LogLevel;
import org.apache.solr.util.TestInjection;
import org.apache.solr.util.TimeOut;
import org.junit.After;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

/** */
@LogLevel(
    "org.apache.solr.cloud.api.collections.ReindexCollectionCmd=DEBUG;org.apache.solr.cloud.OverseerTaskProcessor=DEBUG")
public class ReindexCollectionTest extends SolrCloudTestCase {

  @BeforeClass
  public static void setupCluster() throws Exception {
    configureCluster(2)
        // only *_s
        .addConfig("conf1", configset("cloud-minimal"))
        // every combination of field flags
        .addConfig("conf2", configset("cloud-dynamic"))
        // catch-all * field, indexed+stored
        .addConfig("conf3", configset("cloud-minimal-inplace-updates"))
        .configure();
  }

  private CloudSolrClient solrClient;
  private SolrCloudManager cloudManager;
  private DistribStateManager stateManager;

  @Before
  public void doBefore() {
    ZkController zkController = cluster.getJettySolrRunner(0).getCoreContainer().getZkController();
    cloudManager = zkController.getSolrCloudManager();
    stateManager = cloudManager.getDistribStateManager();
    solrClient =
        new RandomizingCloudSolrClientBuilder(
                Collections.singletonList(zkController.getZkServerAddress()), Optional.empty())
            .build();
  }

  /** Wait for the reindexing state to have the specified value, or fail (can be null). */
  private void waitForReindexingState(String collection, ReindexCollectionCmd.State expected) {

    ReindexCollectionCmd.State lastSeen = null;

    TimeOut timeout = new TimeOut(30, TimeUnit.SECONDS, TimeSource.NANO_TIME);
    while (!timeout.hasTimedOut()) {
      try {
        lastSeen =
            ReindexCollectionCmd.State.get(
                ReindexCollectionCmd.getReindexingState(stateManager, collection)
                    .get(ReindexCollectionCmd.STATE));

        // exit busy waiting if expected state was found
        if (lastSeen == expected) {
          break;
        }
      } catch (Exception e) {
        fail("Unexpected exception checking state of " + collection + ": " + e);
        return;
      }
    }

    assertFalse(
        String.format(Locale.ROOT, "Expecting state %s, last seen was %s", expected, lastSeen),
        timeout.hasTimedOut());
  }

  @After
  public void doAfter() throws Exception {
    cluster.deleteAllCollections(); // deletes aliases too

    if (null != solrClient) {
      solrClient.close();
      solrClient = null;
    }

    TestInjection.reset();
  }

  private static final int NUM_DOCS = 200; // at least two batches, default batchSize=100

  @Test
  public void testBasicReindexing() throws Exception {
    final String sourceCollection = "basicReindexing";

    createCollection(sourceCollection, "conf1", 2, 2);

    indexDocs(
        sourceCollection,
        NUM_DOCS,
        i -> new SolrInputDocument("id", String.valueOf(i), "string_s", String.valueOf(i)));

    final String targetCollection = "basicReindexingTarget";

    CollectionAdminRequest.ReindexCollection req =
        CollectionAdminRequest.reindexCollection(sourceCollection).setTarget(targetCollection);
    CollectionAdminResponse rsp = req.process(solrClient);
    assertNotNull(rsp.toString(), rsp.getResponse().get(ReindexCollectionCmd.REINDEX_STATUS));
    @SuppressWarnings({"unchecked"})
    Map<String, Object> status =
        (Map<String, Object>) rsp.getResponse().get(ReindexCollectionCmd.REINDEX_STATUS);
    assertEquals(
        status.toString(), (long) NUM_DOCS, ((Number) status.get("inputDocs")).longValue());
    assertEquals(
        status.toString(), (long) NUM_DOCS, ((Number) status.get("processedDocs")).longValue());

    waitForState(
        "did not finish copying in time",
        targetCollection,
        coll -> {
          ReindexCollectionCmd.State state =
              ReindexCollectionCmd.State.get(coll.getStr(ReindexCollectionCmd.REINDEXING_STATE));
          return ReindexCollectionCmd.State.FINISHED == state;
        });

    // verify the target docs exist
    QueryResponse queryResponse = solrClient.query(targetCollection, params(CommonParams.Q, "*:*"));
    assertEquals("copied num docs", NUM_DOCS, queryResponse.getResults().getNumFound());
  }

  @Test
  public void testSameTargetReindexing() throws Exception {
    doTestSameTargetReindexing(false, false);
    doTestSameTargetReindexing(false, true);
    doTestSameTargetReindexing(true, false);
    doTestSameTargetReindexing(true, true);
  }

  private void doTestSameTargetReindexing(boolean sourceRemove, boolean followAliases)
      throws Exception {
    final String sourceCollection = "sameTargetReindexing_" + sourceRemove + "_" + followAliases;
    final String targetCollection = sourceCollection;

    createCollection(sourceCollection, "conf1", 2, 2);
    indexDocs(
        sourceCollection,
        NUM_DOCS,
        i -> new SolrInputDocument("id", String.valueOf(i), "string_s", String.valueOf(i)));

    CollectionAdminRequest.ReindexCollection req =
        CollectionAdminRequest.reindexCollection(sourceCollection).setTarget(targetCollection);
    req.setRemoveSource(sourceRemove);
    req.setFollowAliases(followAliases);
    req.process(solrClient);

    String realTargetCollection = null;
    TimeOut timeOut = new TimeOut(30, TimeUnit.SECONDS, cloudManager.getTimeSource());
    String prefix = ReindexCollectionCmd.TARGET_COL_PREFIX + targetCollection;
    while (!timeOut.hasTimedOut()) {
      timeOut.sleep(500);
      for (String name : cloudManager.getClusterState().getCollectionNames()) {
        if (name.startsWith(prefix)) {
          realTargetCollection = name;
          break;
        }
      }
      if (realTargetCollection != null) {
        break;
      }
    }
    assertNotNull("target collection not present after 30s", realTargetCollection);

    waitForState(
        "did not finish copying in time",
        realTargetCollection,
        coll -> {
          ReindexCollectionCmd.State state =
              ReindexCollectionCmd.State.get(coll.getStr(ReindexCollectionCmd.REINDEXING_STATE));
          return ReindexCollectionCmd.State.FINISHED == state;
        });
    ZkStateReader.from(solrClient).aliasesManager.update();
    // verify the target docs exist
    QueryResponse rsp = solrClient.query(targetCollection, params(CommonParams.Q, "*:*"));
    assertEquals("copied num docs", NUM_DOCS, rsp.getResults().getNumFound());
    ClusterState state = solrClient.getClusterState();
    if (sourceRemove) {
      assertFalse("source collection still present", state.hasCollection(sourceCollection));
    }
  }

  @Test
  public void testLossySchema() throws Exception {
    final String sourceCollection = "sourceLossyReindexing";
    final String targetCollection = "targetLossyReindexing";

    createCollection(sourceCollection, "conf2", 2, 2);

    indexDocs(
        sourceCollection,
        NUM_DOCS,
        i ->
            new SolrInputDocument(
                "id", String.valueOf(i),
                "string_s", String.valueOf(i),
                "sind",
                    "this is a test " + i)); // "sind": indexed=true, stored=false, will be lost...

    CollectionAdminRequest.ReindexCollection req =
        CollectionAdminRequest.reindexCollection(sourceCollection)
            .setTarget(targetCollection)
            .setConfigName("conf3");
    req.process(solrClient);

    waitForState(
        "did not finish copying in time",
        targetCollection,
        coll -> {
          ReindexCollectionCmd.State state =
              ReindexCollectionCmd.State.get(coll.getStr(ReindexCollectionCmd.REINDEXING_STATE));
          return ReindexCollectionCmd.State.FINISHED == state;
        });
    // verify the target docs exist
    QueryResponse rsp = solrClient.query(targetCollection, params(CommonParams.Q, "*:*"));
    assertEquals("copied num docs", NUM_DOCS, rsp.getResults().getNumFound());
    for (SolrDocument doc : rsp.getResults()) {
      String id = (String) doc.getFieldValue("id");
      assertEquals(id, doc.getFieldValue("string_s"));
      assertFalse(doc.containsKey("sind")); // lost in translation ...
    }
  }

  @Test
  public void testReshapeReindexing() throws Exception {
    final String sourceCollection = "reshapeReindexing";
    final String targetCollection = "reshapeReindexingTarget";
    createCollection(sourceCollection, "conf1", 2, 2);
    indexDocs(
        sourceCollection,
        NUM_DOCS,
        i ->
            new SolrInputDocument(
                "id", String.valueOf(i),
                "string_s", String.valueOf(i),
                "remove_s", String.valueOf(i)));

    CollectionAdminRequest.ReindexCollection req =
        CollectionAdminRequest.reindexCollection(sourceCollection)
            .setTarget(targetCollection)
            .setCollectionParam(ZkStateReader.NUM_SHARDS_PROP, 3)
            .setCollectionParam(ZkStateReader.REPLICATION_FACTOR, 1)
            .setCollectionParam("router.name", ImplicitDocRouter.NAME)
            .setCollectionParam("shards", "foo,bar,baz")
            .setCollectionParam("fl", "id,string_s")
            .setCollectionParam("q", "id:10*");
    req.process(solrClient);

    waitForState(
        "did not finish copying in time",
        targetCollection,
        coll -> {
          ReindexCollectionCmd.State state =
              ReindexCollectionCmd.State.get(coll.getStr(ReindexCollectionCmd.REINDEXING_STATE));
          return ReindexCollectionCmd.State.FINISHED == state;
        });
    // verify the target docs exist
    QueryResponse rsp = solrClient.query(targetCollection, params(CommonParams.Q, "*:*"));

    // 10 and 100-109
    assertEquals("copied num docs", 11, rsp.getResults().getNumFound());
    // verify the correct fields exist
    for (SolrDocument doc : rsp.getResults()) {
      assertNotNull(doc.getFieldValue("id"));
      assertNotNull(doc.getFieldValue("string_s"));
      assertNull(doc.getFieldValue("remove_s"));
    }

    // check the shape of the new collection
    ClusterState clusterState = solrClient.getClusterState();
    List<String> aliases =
        ZkStateReader.from(solrClient).getAliases().resolveAliases(targetCollection);
    assertFalse(aliases.isEmpty());
    String realTargetCollection = aliases.get(0);
    DocCollection coll = clusterState.getCollection(realTargetCollection);
    assertNotNull(coll);
    assertEquals(3, coll.getSlices().size());
    assertNotNull("foo", coll.getSlice("foo"));
    assertNotNull("bar", coll.getSlice("bar"));
    assertNotNull("baz", coll.getSlice("baz"));
    assertEquals(Integer.valueOf(1), coll.getReplicationFactor());
    assertEquals(ImplicitDocRouter.NAME, coll.getRouter().getName());
  }

  @Test
  public void testFailure() throws Exception {
    final String sourceCollection = "failReindexing";
    final String targetCollection = "failReindexingTarget";
    final String aliasTarget = "failAlias";
    createCollection(sourceCollection, "conf1", 2, 2);
    createCollection(targetCollection, "conf1", 1, 1);
    CollectionAdminRequest.createAlias(aliasTarget, targetCollection).process(solrClient);
    indexDocs(
        sourceCollection,
        NUM_DOCS,
        i ->
            new SolrInputDocument(
                "id", String.valueOf(i),
                "string_s", String.valueOf(i)));

    CollectionAdminRequest.ReindexCollection req =
        CollectionAdminRequest.reindexCollection(sourceCollection).setTarget(targetCollection);
    CollectionAdminResponse rsp = req.process(solrClient);
    assertNotNull(rsp.getResponse().get("error"));
    assertTrue(
        rsp.toString(), rsp.getResponse().get("error").toString().contains("already exists"));

    req = CollectionAdminRequest.reindexCollection(sourceCollection).setTarget(aliasTarget);
    rsp = req.process(solrClient);
    assertNotNull(rsp.getResponse().get("error"));
    assertTrue(
        rsp.toString(), rsp.getResponse().get("error").toString().contains("already exists"));

    CollectionAdminRequest.deleteAlias(aliasTarget).process(solrClient);
    CollectionAdminRequest.deleteCollection(targetCollection).process(solrClient);

    req = CollectionAdminRequest.reindexCollection(sourceCollection).setTarget(targetCollection);

    TestInjection.reindexFailure = "true:100";
    rsp = req.process(solrClient);
    assertNotNull(rsp.getResponse().get("error"));
    assertTrue(
        rsp.toString(), rsp.getResponse().get("error").toString().contains("waiting for daemon"));

    // verify that the target and checkpoint collections don't exist
    cloudManager
        .getClusterState()
        .collectionStream()
        .forEach(
            coll -> {
              assertFalse(
                  coll.getName() + " still exists",
                  coll.getName().startsWith(ReindexCollectionCmd.TARGET_COL_PREFIX));
              assertFalse(
                  coll.getName() + " still exists",
                  coll.getName().startsWith(ReindexCollectionCmd.CHK_COL_PREFIX));
            });
    // verify that the source collection is read-write and has no reindexing flags
    waitForState(
        "collection state is incorrect",
        sourceCollection,
        collectionState ->
            !collectionState.isReadOnly()
                && collectionState.getStr(ReindexCollectionCmd.REINDEXING_STATE) == null);
    waitForReindexingState(sourceCollection, null);
  }

  @Test
  @SuppressWarnings({"unchecked"})
  public void testAbort() throws Exception {
    final String sourceCollection = "abortReindexing";
    final String targetCollection = "abortReindexingTarget";
    createCollection(sourceCollection, "conf1", 2, 1);

    TestInjection.reindexLatch = new CountDownLatch(1);
    CollectionAdminRequest.ReindexCollection req =
        CollectionAdminRequest.reindexCollection(sourceCollection).setTarget(targetCollection);
    String asyncId = req.processAsync(solrClient);
    // wait for the source collection to be put in readOnly mode
    waitForState(
        "source collection didn't become readOnly", sourceCollection, DocCollection::isReadOnly);

    req = CollectionAdminRequest.reindexCollection(sourceCollection);
    req.setCommand("abort");
    CollectionAdminResponse rsp = req.process(solrClient);
    Map<String, Object> status =
        (Map<String, Object>) rsp.getResponse().get(ReindexCollectionCmd.REINDEX_STATUS);
    assertNotNull(rsp.toString(), status);
    assertEquals(status.toString(), "aborting", status.get("state"));

    waitForState("incorrect collection state", sourceCollection, DocCollection::isReadOnly);
    waitForReindexingState(sourceCollection, ReindexCollectionCmd.State.ABORTED);

    // verify status
    req.setCommand("status");
    rsp = req.process(solrClient);
    status = (Map<String, Object>) rsp.getResponse().get(ReindexCollectionCmd.REINDEX_STATUS);
    assertNotNull(rsp.toString(), status);
    assertEquals(status.toString(), "aborted", status.get("state"));
    // let the process continue
    TestInjection.reindexLatch.countDown();
    waitForState(
        "source collection is in wrong state",
        sourceCollection,
        docCollection -> !docCollection.isReadOnly());
    waitForReindexingState(sourceCollection, null);

    // verify the response
    rsp = CollectionAdminRequest.requestStatus(asyncId).process(solrClient);
    int retries = 0;
    while (retries < 5
        && ((CollectionAdminRequest.RequestStatusResponse) rsp).getRequestStatus()
            == RequestStatusState.RUNNING) {
      // There is a chance the async status hasn't been updated yet when we retrieve it the first
      // time. Just wait and retry a couple of times
      Thread.sleep(100L);
      rsp = CollectionAdminRequest.requestStatus(asyncId).process(solrClient);
      retries++;
    }
    status = (Map<String, Object>) rsp.getResponse().get(ReindexCollectionCmd.REINDEX_STATUS);
    assertNotNull(rsp.toString(), status);
    assertEquals(status.toString(), "aborted", status.get("state"));
  }

  private void createCollection(String name, String config, int numShards, int numReplicas)
      throws Exception {
    CollectionAdminRequest.createCollection(name, config, numShards, numReplicas)
        .process(solrClient);

    cluster.waitForActiveCollection(name, numShards, numShards * numReplicas);
  }

  private void indexDocs(
      String collection, int numDocs, Function<Integer, SolrInputDocument> generator)
      throws Exception {
    List<SolrInputDocument> docs = new ArrayList<>();
    for (int i = 0; i < numDocs; i++) {
      docs.add(generator.apply(i));
    }
    solrClient.add(collection, docs);
    solrClient.commit(collection);
    // verify the docs exist
    QueryResponse rsp = solrClient.query(collection, params(CommonParams.Q, "*:*"));

    assertEquals("num docs", NUM_DOCS, rsp.getResults().getNumFound());
  }
}
